
#include "lib.h"

#include <stdlib.h>
#include <memory.h>

#include "serverfile.h"
#include "serializer.h"
#include "md4.h"

CServerFile::CServerFile()
{
	if(!InitState())
	{
		// throw an exception!
	}
}

bool8 
CServerFile::InitState()
{
	try
	{
		m_obuf			= (uint8 *)0;
		m_num_blocks 	= 0;

		m_buffer 		= (uint8 *)0;
		m_bufsize 		= 0;

		m_block_map 	= (uint32 *)0;
		m_weak_checksums = (uint32 *)0;
		m_md4_checksums = (t_md4 *)0;

		m_last_byte 	= 0;

		m_a = m_b 		= 0;

		m_block_id 		= -1;

		// initially we need to calculate m_a and m_b directly 
		m_can_roll 		= false;

		memset(m_hash_table, 0, sizeof(uint32 *) * HASH_SIZE);

		return true;
	}
	catch(...)
	{
		// access violation?
		return false;
	}
}

bool8
CServerFile::Done()
{
	try
	{
		delete [] m_weak_checksums;
		m_weak_checksums = (uint32 *)0;

		delete [] m_md4_checksums;
		m_md4_checksums = (t_md4 *)0;

		delete [] m_block_map;
		m_block_map 	= (uint32 *)0;

		return true;
	}
	catch(...)
	{
		// access violation?
		return false;
	}
}

CServerFile::~CServerFile()
{
	Done();
}

bool8 
CServerFile::SetupHashTable()
{
    if (!m_num_blocks)
    {
        return true;
    }
	else if(m_weak_checksums)
	{
		memset(m_hash_table, 0, sizeof(uint32 *) * HASH_SIZE);

		uint32 i = 0;
		uint16 current_hash;

		while(i < m_num_blocks)
		{
			current_hash = (uint16)(m_weak_checksums[i] >> 16);
			m_hash_table[current_hash] = &m_weak_checksums[i];

			while((i < m_num_blocks) && 
				((uint16)(m_weak_checksums[i] >> 16) == current_hash)) i++;
		}

		return true;
	}
	else
	{
		return false;
	}
}

bool8 
CServerFile::SortChecksums()
{
    if (!m_num_blocks)
    {
        return true;
    }
	else if(m_weak_checksums && m_md4_checksums)
	{
		// ok this is sort of cheesy, but it lets us use the qsort function.
		//  since we want to sort the weak checksums, but we need the indexes
		//  of the weak and the strong checksums to coincide, we can't just 
		//  srot the weaks as it'll get the two out of sync.  so we make a temp
		//  array with a pointer to each and then sort those.

		t_weak_and_strong *temp = new t_weak_and_strong[m_num_blocks];
		if(!temp)
		{
			ASSERT(("out of memory!", 0));
			return false;
		}

		for(uint32 i = 0; i < m_num_blocks; i++)
		{
			temp[i].weak = &m_weak_checksums[i];
			temp[i].strong = &m_md4_checksums[i];
		}

		qsort((void *)temp, m_num_blocks, 
			sizeof(t_weak_and_strong), checksum_compare);

		// and now rearrange then into a new array and blow out the old one.
		uint32 *new_weak;
		t_md4  *new_strong;

		new_weak = new uint32[m_num_blocks];
		if(!new_weak)
		{
			delete [] temp;
			return false;
		}

		new_strong = new t_md4[m_num_blocks];
		if(!new_strong)
		{
			delete [] new_weak;
			delete [] temp;
			return false;
		}

		for(i = 0; i < m_num_blocks; i++)
		{
			new_weak[i] = *temp[i].weak;
			new_strong[i][0]  = (*(temp[i].strong))[0];
			new_strong[i][1]  = (*(temp[i].strong))[1];
			new_strong[i][2]  = (*(temp[i].strong))[2];
			new_strong[i][3]  = (*(temp[i].strong))[3];
		}

		// when we send info back to the client, we must reference blocks
		//  in their original file, so we must create a mapping from the new
		//  organization to the old.
		if(!CreateBlockMap(new_weak))
		{
			delete [] new_weak;
			delete [] new_strong;
			delete [] temp;

			return false;
		}

		delete [] m_weak_checksums;
		m_weak_checksums = (uint32 *)0;

		delete [] m_md4_checksums;
		m_md4_checksums = (t_md4 *)0;

		m_weak_checksums = new_weak;
		m_md4_checksums = new_strong;

		delete [] temp;

		return true;
	}
	else
	{
		return false;
	}
}

// yes, i COULD have used an ordered list for all of this, but i 
//  really want it to be fast!  and it's not fast right now, i know,
//  but i want it to have POTENTIAL! :P
bool8 
CServerFile::CreateBlockMap(uint32 *new_weak_checksums)
{
	if(m_num_blocks)
	{
		struct t_map_entry
		{
			uint32 new_block_idx;
			uint32 old_block_idx;
		};

		t_map_entry *block_map;
		block_map = new t_map_entry[m_num_blocks];
		if(!block_map)
		{
			// out of memory!
			return false;
		}

		m_block_map = new uint32[m_num_blocks];
		if(!m_block_map)
		{
			delete [] block_map;
			return false;
		}

		for(uint32 i = 0; i < m_num_blocks; i++)
		{
			block_map[i].old_block_idx = i;
			block_map[i].new_block_idx = 
				FindWeakCS(m_weak_checksums[i], new_weak_checksums);
		}

		// and invert the array for quick indexing
		for(i = 0; i < m_num_blocks; i++)
		{
			m_block_map[block_map[i].new_block_idx] = 
				block_map[i].old_block_idx;
		}

		delete [] block_map;
	}

	return true;
}

uint32 
CServerFile::FindWeakCS(uint32 weakCS, uint32 *checksums)
{
	for(uint32 i = 0; i < m_num_blocks; i++)
	{
		if(checksums[i] == weakCS)
			return i;
	}

	ASSERT(("couldn't find the checksum!\n", 0));
	return 0;
}

int 
CServerFile::checksum_compare(const void *c1, const void *c2)
{
	t_weak_and_strong *v1 = (t_weak_and_strong *)c1;
	t_weak_and_strong *v2 = (t_weak_and_strong *)c2;

	if(*v1->weak < *v2->weak) 
		return -1;
	else if(*v1->weak > *v2->weak) 
		return 1;
	else 
		return 0;
}

bool8
CServerFile::FindStrongChecksum(t_md4 md4, uint8 *buffer)
{
	// the MD4 algorithm uses blocks of length 512 BITS.
	//  512 BITS == 64 BYTES, so compensate appropriately
	//  and 512 BYTES == 8 64 bit chunks
	if(buffer)
	{
		MDstruct md4struct;
		MDbegin(&md4struct);

		int32 count = (BLOCK_SIZE / 64);
		for(int32 i = 0; i < count; i++)
		{
			MDupdate(&md4struct, buffer, 512);
		}
		MDupdate(&md4struct, buffer, 0);

		md4[0] = md4struct.buffer[0];
		md4[1] = md4struct.buffer[1];
		md4[2] = md4struct.buffer[2];
		md4[3] = md4struct.buffer[3];

		return true;
	}
	else
	{
		return false;
	}
}

uint32 
CServerFile::FindRollingChecksum(uint8 *buffer)
{
	// just a sanity check
	if(!buffer)
		return false;

	try
	{

		// find the rolling checksum
		if(!m_can_roll)
		{
			// calculate the weak checksum from the entire block
			// we don't have a previous checksum, so we just calculate it...
			uint32 i;
			uint8 *buf = m_buffer;

			m_a = m_b = 0;

			for(i = 0; (int32)i <= BLOCK_SIZE - 4; i += 4)
			{
				m_b += 
					((m_a + buf[i]) << 2) 	+
					3 * buf[i + 1] 			+
					(buf[i + 2] << 1)		+
					buf[i + 3];

				m_a += 
					buf[i] 		+ 
					buf[i + 1] 	+ 
					buf[i + 2] 	+ 
					buf[i + 3];
			}

			while(i < (uint32)BLOCK_SIZE)
			{
				m_a += buf[i];
				m_b += m_a;
				i++;
			}
		}
		else
		{
			m_a -= m_last_byte;
			m_b -= BLOCK_SIZE * m_last_byte;

			m_a += buffer[BLOCK_SIZE - 1];
			m_b += m_a;
		}

		m_last_byte = buffer[0];
		m_can_roll = true;

		return (uint32)((m_a & 0xFFFF) | (m_b << 16));
	}
	catch(...)
	{
		// good as any, i guess...
		return 0;
	}
}

bool8 
CServerFile::FoundBlock(uint32 found_block)
{
	m_block_id = found_block;
	m_can_roll = false;
	return true;
}

bool8 
CServerFile::MissedBlock(uint8 missed_byte)
{
	m_obuf += CSerializer::Wire(m_obuf, missed_byte);
	return true;
}

int32 
CServerFile::Digest(uint32 amount, uint8 *obuf, int32 obufsize)
{
	try
	{
        if (!m_num_blocks)
        {
            uint32 readAmt = amount;
            if (m_bufsize < readAmt)
                readAmt = m_bufsize;
            if ((uint32)obufsize < readAmt)
                readAmt = (uint32)obufsize;
			m_obuf = obuf;
            memcpy(m_obuf, m_buffer, readAmt);
            m_obuf += readAmt;
            return (int32)readAmt;
        }
		else if(m_weak_checksums && m_md4_checksums)
		{
			uint32 *current_hash = (uint32 *)0;
			uint32 	weak_value;
			t_md4	md4;
			bool8	found;

			// start looking at the block represented beween m_buffer[slider]
			//  and m_buffer[slider + BLOCK_SIZE].
			uint32 slider = 0;
			uint32 dstslider;

			if(amount > BLOCK_SIZE)
			{
				dstslider = amount - BLOCK_SIZE;
			}
			else
			{
				// if the amount passed into digest is less than an entire
				//  block, then this must be the last block of the file.
				dstslider = amount;
			}

			m_obuf = obuf;

			// we continue digesting until there's nothing more to digest or
			//  our out buffer is full
			while((slider < dstslider) && 
				((m_obuf - obuf) < (obufsize - BLOCK_HEADERSIZE - 1)))
			{
				found = false;
				m_block_id = -1;
				weak_value = FindRollingChecksum(m_buffer + slider);

				// see if we have it hashed.
				//  this is the first level check
				if(m_hash_table[(weak_value >> 16)])
				{
					// we do have a hash entry for it, and all the weak 
					//  checksums are sorted, so look to see if we have the
					//  full weak checksum here
					//  this is the second level of comparison
					current_hash = m_hash_table[(weak_value >> 16)];

					while((current_hash < m_weak_checksums + m_num_blocks) &&
						(*current_hash < weak_value))
					{
						current_hash++;
					}

					if((current_hash < m_weak_checksums + m_num_blocks) && 
						(*current_hash == weak_value))
					{
						// we do, so we perform the third and final check.
						//  we check the md4 fingerprint of both and see if they
						//  correspond
						// find tmd4
						FindStrongChecksum(md4, m_buffer + slider);

						// compare the two
						if(!memcmp(md4, 
							m_md4_checksums[current_hash - m_weak_checksums], 
							sizeof(uint32) * 4))
						{
							// and we have a match!
							found = true;
						}
					}
				}

				if(found)
				{
					if(!FoundBlock(m_block_map[current_hash - m_weak_checksums]))
					{
						return -1;
					}
					slider += BLOCK_SIZE;

					if(slider > amount) 
						slider = amount;

					return slider;
				}
				else
				{
					if(!MissedBlock(m_buffer[slider]))
					{
						return -1;
					}
					slider++;
				}

			}

			return slider;
		}
		else
		{
			return -1;
		}
	}
	catch(...)
	{
		return -1;
	}
}

bool8 
CServerFile::DoneChecksums()
{
    if (!m_num_blocks)
    {
        return true;
    }
	else if(m_weak_checksums && m_md4_checksums)
	{
		if(!SortChecksums())
		{
			return false;
		}

		if(!SetupHashTable())
		{
			return false;
		}

		return true;
	}
	else
	{
		return false;
	}
}




